from flask import Flask, render_template, request, redirect, url_for, flash, jsonify
from flask_sqlalchemy import SQLAlchemy
from datetime import datetime
from sqlalchemy import extract
import plotly.graph_objects as go
import pandas as pd
from flask_login import LoginManager, UserMixin, login_user, login_required, logout_user, current_user

app = Flask(__name__)
app.config['SECRET_KEY'] = 'your-secret-key'
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///finance.db'
db = SQLAlchemy(app)

# 启用SQLite外键支持
from sqlalchemy import event
from sqlalchemy.engine import Engine

@event.listens_for(Engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("PRAGMA foreign_keys=ON")
    cursor.close()
login_manager = LoginManager(app)
login_manager.login_view = 'login'

# User model

@app.route('/delete_account', methods=['POST'])
@login_required
def delete_account():
    user = User.query.get(current_user.id)
    logout_user()
    db.session.delete(user)
    db.session.commit()
    flash('账户已注销，感谢使用！')
    return redirect(url_for('register'))


@app.route('/api/savings_goal/<int:goal_id>')
@login_required
def api_get_savings_goal(goal_id):
    goal = SavingsGoal.query.filter_by(id=goal_id, user_id=current_user.id).first()
    if not goal:
        return jsonify({'success': False, 'message': '目标不存在'}), 404
    return jsonify({
        'id': goal.id,
        'name': goal.name,
        'target_amount': goal.target_amount,
        'current_amount': goal.current_amount,
        'target_date': goal.target_date.strftime('%Y-%m-%d')
    })

class User(UserMixin, db.Model):
    __tablename__ = 'user'
    __table_args__ = {'extend_existing': True}
    id = db.Column(db.Integer, primary_key=True)
    username = db.Column(db.String(80), unique=True, nullable=False)
    password = db.Column(db.String(120), nullable=False)
    current_balance = db.Column(db.Float, nullable=True, default=0.0)
    categories = db.relationship('Category', backref='user', lazy=True, cascade="all, delete-orphan")
    transactions = db.relationship('Transaction', backref='user', lazy=True, cascade="all, delete-orphan")
    savings_goals = db.relationship('SavingsGoal', backref='user', lazy=True, cascade="all, delete-orphan")

class Category(db.Model):
    __tablename__ = 'categories'
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id', ondelete="CASCADE"), nullable=False)
    name = db.Column(db.String(50), nullable=False)
    type = db.Column(db.String(10), nullable=False)  # 'income' or 'expense'
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

    def __repr__(self):
        return f'<Category {self.name}>'

class Transaction(db.Model):
    __tablename__ = 'transactions'
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id', ondelete="CASCADE"), nullable=False)
    amount = db.Column(db.Float, nullable=False)
    category_id = db.Column(db.Integer, db.ForeignKey('categories.id', ondelete="CASCADE"), nullable=False)
    date = db.Column(db.DateTime, nullable=False, default=datetime.utcnow)
    description = db.Column(db.String(200))
    type = db.Column(db.String(10), nullable=False)  # 'income' or 'expense'
    category = db.relationship('Category', backref='transactions', lazy=True)

class SavingsGoal(db.Model):
    id = db.Column(db.Integer, primary_key=True)
    user_id = db.Column(db.Integer, db.ForeignKey('user.id', ondelete="CASCADE"), nullable=False)
    name = db.Column(db.String(100), nullable=False)
    target_amount = db.Column(db.Float, nullable=False)
    current_amount = db.Column(db.Float, default=0)
    target_date = db.Column(db.DateTime, nullable=False)

@login_manager.user_loader
def load_user(user_id):
    return User.query.get(int(user_id))

@app.route('/')
def index():
    return redirect(url_for('login'))

@app.route('/login', methods=['GET', 'POST'])
def login():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        user = User.query.filter_by(username=username).first()
        if user and user.password == password:  # In production, use proper password hashing
            login_user(user)
            return redirect(url_for('dashboard'))
        flash('Invalid username or password')
    return render_template('login.html')

@app.route('/register', methods=['GET', 'POST'])
def register():
    if request.method == 'POST':
        username = request.form.get('username')
        password = request.form.get('password')
        confirm_password = request.form.get('confirm_password')
        
        if password != confirm_password:
            flash('Passwords do not match')
            return redirect(url_for('register'))
            
        existing_user = User.query.filter_by(username=username).first()
        if existing_user:
            flash('Username already exists')
            return redirect(url_for('register'))
            
        new_user = User(username=username, password=password)
        db.session.add(new_user)
        db.session.commit()
        
        flash('Registration successful! Please login.')
        return redirect(url_for('login'))
    
    return render_template('register.html')

@app.route('/logout')
@login_required
def logout():
    logout_user()
    return redirect(url_for('login'))

@app.route('/account_info')
@login_required
def account_info():
    return render_template('account_info.html', user=current_user)

@app.route('/change_password', methods=['GET', 'POST'])
@login_required
def change_password():
    if request.method == 'POST':
        old_password = request.form.get('old_password')
        new_password = request.form.get('new_password')
        confirm_password = request.form.get('confirm_password')
        if old_password != current_user.password:
            flash('原密码错误')
            return redirect(url_for('change_password'))
        if new_password != confirm_password:
            flash('两次输入的新密码不一致')
            return redirect(url_for('change_password'))
        current_user.password = new_password
        db.session.commit()
        flash('密码修改成功，请重新登录')
        return redirect(url_for('logout'))
    return render_template('change_password.html')


@app.route('/dashboard')
@login_required
def dashboard():
    # Get current month's transactions
    current_month = datetime.now().month
    current_year = datetime.now().year
    monthly_transactions = Transaction.query.filter(
        Transaction.user_id == current_user.id,
        extract('month', Transaction.date) == current_month,
        extract('year', Transaction.date) == current_year
    ).all()
    
    # Calculate total income and expenses
    total_income = sum(t.amount for t in monthly_transactions if t.type == 'income')
    total_expenses = sum(t.amount for t in monthly_transactions if t.type == 'expense')
    
    # Create pie chart for expense categories
    categories = [t.category.name if t.category else '未知分类' for t in monthly_transactions if t.type == 'expense']
    amounts = [t.amount for t in monthly_transactions if t.type == 'expense']
    
    fig = go.Figure(data=[go.Pie(labels=categories, values=amounts)])
    fig.update_layout(title_text='本月支出分类分布')
    pie_chart = fig.to_html(full_html=False)

    # 新增：本年每月收入和支出折线图
    current_year = datetime.now().year
    months = list(range(1, 13))
    month_labels = [f"{m}月" for m in months]
    # 查询本年所有交易
    year_transactions = Transaction.query.filter(
        Transaction.user_id == current_user.id,
        extract('year', Transaction.date) == current_year
    ).all()
    # 统计每月收入和支出
    monthly_income = [sum(t.amount for t in year_transactions if t.type == 'income' and t.date.month == m) for m in months]
    monthly_expense = [sum(t.amount for t in year_transactions if t.type == 'expense' and t.date.month == m) for m in months]

    # 收入折线图
    income_fig = go.Figure()
    income_fig.add_trace(go.Scatter(x=month_labels, y=monthly_income, mode='lines+markers', name='收入', line=dict(color='green')))
    income_fig.update_layout(title_text='本年每月收入趋势', xaxis_title='月份', yaxis_title='收入金额')
    income_line_chart = income_fig.to_html(full_html=False)

    # 支出折线图
    expense_fig = go.Figure()
    expense_fig.add_trace(go.Scatter(x=month_labels, y=monthly_expense, mode='lines+markers', name='支出', line=dict(color='red')))
    expense_fig.update_layout(title_text='本年每月支出趋势', xaxis_title='月份', yaxis_title='支出金额')
    expense_line_chart = expense_fig.to_html(full_html=False)
    
    # Get savings goals
    savings_goals = SavingsGoal.query.filter_by(user_id=current_user.id).all()
    
    # Get all categories for the quick transaction form
    categories = Category.query.filter_by(user_id=current_user.id).all()
    
    # Get current balance
    current_balance = current_user.current_balance
    
    return render_template('dashboard.html',
                         total_income=total_income,
                         total_expenses=total_expenses,
                         pie_chart=pie_chart,
                         income_line_chart=income_line_chart,
                         expense_line_chart=expense_line_chart,
                         savings_goals=savings_goals,
                         categories=categories,
                         current_balance=current_balance,
                         current_date=datetime.now().strftime('%Y-%m-%d'))


@app.route('/savings_goals')
@login_required
def savings_goals():
    goals = SavingsGoal.query.filter_by(user_id=current_user.id).all()
    return render_template('savings_goals.html', goals=goals)

@app.route('/deposit_to_goal', methods=['POST'])
@login_required
def deposit_to_goal():
    try:
        goal_id = int(request.form['goal_id'])
        amount = float(request.form['amount'])
        
        # 验证目标存在且属于当前用户
        goal = SavingsGoal.query.filter_by(id=goal_id, user_id=current_user.id).first()
        if not goal:
            return jsonify({'success': False, 'message': '目标不存在'}), 404
            
        # 检查余额是否足够
        if current_user.current_balance < amount:
            return jsonify({'success': False, 'message': '余额不足'}), 400
            
        # 更新目标金额和用户余额
        goal.current_amount += amount
        current_user.current_balance -= amount
        
        # 查找或创建“储蓄”分类（类型可设为“储蓄”或“支出”）
        savings_category = Category.query.filter_by(user_id=current_user.id, name="储蓄").first()
        if not savings_category:
            savings_category = Category(user_id=current_user.id, name="储蓄", type="储蓄")
            db.session.add(savings_category)
            db.session.commit()  # 立即提交以获得ID
        # 创建交易记录（类型为“储蓄”）
        transaction = Transaction(
            user_id=current_user.id,
            amount=amount,
            category_id=savings_category.id,  # 指定“储蓄”分类ID
            description=f'存入储蓄目标: {goal.name}',
            type='savings',  # 新增储蓄类型
            date=datetime.now()
        )
        
        db.session.add(transaction)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'message': '存入成功',
            'new_balance': current_user.current_balance,
            'new_goal_amount': goal.current_amount
        })
    except ValueError:
        return jsonify({'success': False, 'message': '金额必须是数字'}), 400
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': f'操作失败: {str(e)}'}), 500

@app.route('/withdraw_from_goal', methods=['POST'])
@login_required
def withdraw_from_goal():
    try:
        goal_id = int(request.form['goal_id'])
        amount = float(request.form['amount'])
        
        # 验证目标存在且属于当前用户
        goal = SavingsGoal.query.filter_by(id=goal_id, user_id=current_user.id).first()
        if not goal:
            return jsonify({'success': False, 'message': '目标不存在'}), 404
            
        # 检查目标金额是否足够
        if goal.current_amount < amount:
            return jsonify({'success': False, 'message': '目标金额不足'}), 400
            
        # 更新目标金额和用户余额
        goal.current_amount -= amount
        current_user.current_balance += amount
        
        # 查找或创建“储蓄”分类（类型可设为“储蓄”或“支出”）
        savings_category = Category.query.filter_by(user_id=current_user.id, name="储蓄").first()
        if not savings_category:
            savings_category = Category(user_id=current_user.id, name="储蓄", type="储蓄")
            db.session.add(savings_category)
            db.session.commit()  # 立即提交以获得ID
        # 创建交易记录（类型为“取款”）
        transaction = Transaction(
            user_id=current_user.id,
            amount=amount,
            category_id=savings_category.id,  # 指定“储蓄”分类ID
            description=f'从储蓄目标取出: {goal.name}',
            type='withdraw',  # 新增取款类型
            date=datetime.now()
        )
        
        db.session.add(transaction)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'message': '取出成功',
            'new_balance': current_user.current_balance,
            'new_goal_amount': goal.current_amount
        })
    except ValueError:
        return jsonify({'success': False, 'message': '金额必须是数字'}), 400
    except Exception as e:
        db.session.rollback()
        return jsonify({'success': False, 'message': f'操作失败: {str(e)}'}), 500

@app.route('/categories')
@login_required
def categories():
    categories = Category.query.filter_by(user_id=current_user.id).all()
    return render_template('categories.html', categories=categories)



@app.route('/transaction_details')
@login_required
def transaction_details():
    # 查询所有收入和支出，按日期倒序
    expense_transactions = Transaction.query.filter_by(user_id=current_user.id, type='expense').order_by(Transaction.date.desc()).all()
    income_transactions = Transaction.query.filter_by(user_id=current_user.id, type='income').order_by(Transaction.date.desc()).all()
    return render_template('transaction_details.html', 
                         expense_transactions=expense_transactions,
                         income_transactions=income_transactions)

@app.route('/api/categories')
@login_required
def get_categories():
    transaction_type = request.args.get('type', 'expense')
    categories = Category.query.filter_by(
        user_id=current_user.id,
        type=transaction_type
    ).all()
    return jsonify([{
        'id': category.id,
        'name': category.name
    } for category in categories])

@app.route('/api/transactions')
@login_required
def get_transactions():
    page = int(request.args.get('page', 1))
    size = int(request.args.get('size', 10))
    
    # 获取所有交易记录
    query = Transaction.query.filter_by(user_id=current_user.id)
    total = query.count()
    
    # 分页查询
    transactions = query.order_by(Transaction.date.desc())\
        .paginate(page=page, per_page=size, error_out=False)
    
    # 转换为 JSON 格式
    result = {
        'transactions': [
            {
                'id': t.id,
                'amount': t.amount,
                'type': t.type,
                'category': {
                    'id': t.category_id,
                    'name': t.category.name
                },
                'date': t.date.isoformat(),
                'description': t.description
            }
            for t in transactions.items
        ],
        'totalPages': (total + size - 1) // size,
        'currentPage': page
    }
    
    return jsonify(result)

@app.route('/reset_account', methods=['POST'])
@login_required
def reset_account():
    try:
        # 删除所有相关记录
        # 注意：我们需要先删除依赖关系的表，然后再删除用户表
        Transaction.query.filter_by(user_id=current_user.id).delete()
        Category.query.filter_by(user_id=current_user.id).delete()
        SavingsGoal.query.filter_by(user_id=current_user.id).delete()
        
        # 更新用户余额
        current_user.current_balance = 0.0
        
        db.session.commit()
        
        return jsonify({
            'success': True,
            'message': '账户已重置成功'
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'success': False,
            'message': f'重置账户时出错: {str(e)}'
        }), 500

@app.route('/add_category', methods=['POST'])
@login_required
def add_category():
    name = request.form.get('name')
    type = request.form.get('type')
    
    # 检查是否已存在同名分类
    existing_category = Category.query.filter_by(
        user_id=current_user.id,
        name=name,
        type=type
    ).first()
    
    if existing_category:
        flash('该分类已存在')
        return redirect(url_for('categories'))
    
    new_category = Category(
        user_id=current_user.id,
        name=name,
        type=type
    )
    db.session.add(new_category)
    db.session.commit()
    
    flash('分类已添加')
    return redirect(url_for('categories'))

@app.route('/category/<int:category_id>/delete', methods=['POST'])
@login_required
def delete_category(category_id):
    try:
        category = Category.query.get_or_404(category_id)
        if category.user_id != current_user.id:
            return jsonify({
                'success': False,
                'message': '没有权限删除该分类'
            }), 403
        
        # 检查是否有交易使用该分类
        if category.transactions:
            return jsonify({
                'success': False,
                'message': '该分类正在被使用，无法删除'
            }), 400
        
        db.session.delete(category)
        db.session.commit()
        return jsonify({
            'success': True,
            'message': '分类删除成功'
        })
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'success': False,
            'message': f'删除过程中发生错误: {str(e)}'
        }), 500
    
    return jsonify({'success': True})

@app.route('/save_transaction', methods=['POST'])
@login_required
def save_transaction():
    try:
        # 获取数据
        amount = float(request.form['amount'])
        category_id = int(request.form['category_id'])
        description = request.form.get('description', '')
        date = datetime.strptime(request.form['date'], '%Y-%m-%d')
        type = request.form['type']
        
        # 验证分类是否存在
        category = Category.query.get_or_404(category_id)
        if category.user_id != current_user.id:
            return jsonify({
                'success': False,
                'message': '无法使用其他用户的分类'
            }), 403
        
        # 创建新交易
        transaction = Transaction(
            user_id=current_user.id,
            amount=amount,
            category_id=category_id,
            description=description,
            type=type,
            date=date
        )
        
        # 更新用户余额
        if type == 'income':
            current_user.current_balance += amount
        elif type == 'expense':
            current_user.current_balance -= amount
            if current_user.current_balance < 0:
                return jsonify({
                    'success': False,
                    'message': '余额不足'
                }), 400
        
        db.session.add(transaction)
        db.session.commit()
        
        return jsonify({
            'success': True,
            'message': '交易保存成功',
            'current_balance': current_user.current_balance
        })
    except ValueError:
        return jsonify({
            'success': False,
            'message': '金额必须是数字'
        }), 400
    except Exception as e:
        db.session.rollback()
        return jsonify({
            'success': False,
            'message': f'保存交易时发生错误: {str(e)}'
        }), 500

@app.route('/quick_transaction', methods=['POST'])
@login_required
def quick_transaction():
    try:
        # 获取表单数据
        amount = float(request.form['amount'])
        category_id = int(request.form['category_id'])
        description = request.form.get('description', '')
        date = datetime.strptime(request.form['date'], '%Y-%m-%d')
        type = request.form['type']
        
        # 验证分类是否存在
        category = Category.query.get_or_404(category_id)
        if category.user_id != current_user.id:
            flash('无法使用其他用户的分类')
            return redirect(url_for('dashboard'))
        
        # 创建新交易
        transaction = Transaction(
            user_id=current_user.id,
            amount=amount,
            category_id=category_id,
            description=description,
            type=type,
            date=date
        )
        
        # 更新用户余额
        if type == 'income':
            current_user.current_balance += amount
        elif type == 'expense':
            current_user.current_balance -= amount
            if current_user.current_balance < 0:
                flash('余额不足')
                return redirect(url_for('dashboard'))
        
        db.session.add(transaction)
        db.session.commit()
        # 判断请求类型，支持AJAX和普通跳转
        if request.headers.get('X-Requested-With') == 'XMLHttpRequest':
            return jsonify({'success': True, 'message': '交易保存成功', 'current_balance': current_user.current_balance})
        else:
            flash('交易保存成功')
            return redirect(url_for('dashboard'))
    except ValueError:
        flash('金额必须是数字')
        return redirect(url_for('dashboard'))
    except Exception as e:
        db.session.rollback()
        flash(f'保存交易时发生错误: {str(e)}')
        return redirect(url_for('dashboard'))

@app.route('/add_savings_goal', methods=['POST'])
@login_required
def add_savings_goal():
    name = request.form.get('name')
    target_amount = float(request.form.get('target_amount'))
    target_date = datetime.strptime(request.form.get('target_date'), '%Y-%m-%d')
    
    new_goal = SavingsGoal(
        user_id=current_user.id,
        name=name,
        target_amount=target_amount,
        target_date=target_date
    )
    db.session.add(new_goal)
    db.session.commit()
    
    flash('储蓄目标已添加')
    return redirect(url_for('savings_goals'))
    name = request.form.get('name')
    target_amount = float(request.form.get('target_amount'))
    target_date = datetime.strptime(request.form.get('target_date'), '%Y-%m-%d')
    
    new_goal = SavingsGoal(
        user_id=current_user.id,
        name=name,
        target_amount=target_amount,
        target_date=target_date
    )
    db.session.add(new_goal)
    db.session.commit()
    
    flash('储蓄目标已添加')
    return redirect(url_for('savings_goals'))
    name = request.form.get('name')
    target_amount = float(request.form.get('target_amount'))
    target_date = datetime.strptime(request.form.get('target_date'), '%Y-%m-%d')
    
    new_goal = SavingsGoal(
        user_id=current_user.id,
        name=name,
        target_amount=target_amount,
        target_date=target_date
    )
    db.session.add(new_goal)
    db.session.commit()
    
    flash('储蓄目标已添加')
    return redirect(url_for('savings_goals'))
    name = request.form.get('name')
    target_amount = float(request.form.get('target_amount'))
    target_date = datetime.strptime(request.form.get('target_date'), '%Y-%m-%d')
    
    new_goal = SavingsGoal(
        user_id=current_user.id,
        name=name,
        target_amount=target_amount,
        target_date=target_date
    )
    db.session.add(new_goal)

if __name__ == '__main__':
    # 只在第一次运行时初始化数据库
    from config import is_database_initialized, mark_database_initialized
    
    if not is_database_initialized():
        print('初始化数据库...')
        with app.app_context():
            # 删除所有表
            print('正在删除所有表...')
            db.drop_all()
            print('所有表已删除')
            
            # 创建数据库表
            print('正在创建数据库表...')
            db.create_all()
            print('数据库表创建完成')
            
            # 检查表结构
            print('检查表结构...')
            user_table = db.Table('user', db.metadata)
            print(f'表列: {[c.name for c in user_table.columns]}')
            
            # 添加默认用户
            admin = User(
                username='admin',
                password='admin',
                current_balance=0.0  # 请在生产环境中使用安全的密码哈希
            )
            db.session.add(admin)
            try:
                db.session.commit()
                print('默认用户创建成功')
                
                # 验证用户数据
                user = User.query.first()
                print(f'用户数据: {user.__dict__}')
                
                # 如果 current_balance 不存在，添加它
                if not hasattr(user, 'current_balance'):
                    print('添加 current_balance 字段...')
                    user.current_balance = 0.0
                    db.session.commit()
                    print('current_balance 字段添加成功')
                
                # 标记数据库已初始化
                mark_database_initialized()
            except Exception as e:
                print(f'创建用户时出错: {e}')
                db.session.rollback()
                # 如果创建用户失败，检查表结构
                try:
                    # 检查 user 表是否存在 current_balance 列
                    user_table = db.Table('user', db.metadata)
                    if 'current_balance' not in [c.name for c in user_table.columns]:
                        print('正在添加 current_balance 列...')
                        # 使用 Alembic 创建迁移
                        from alembic.config import Config
                        from alembic import command
                        
                        alembic_cfg = Config()
                        alembic_cfg.set_main_option('script_location', 'migrations')
                        alembic_cfg.set_main_option('sqlalchemy.url', 'sqlite:///finance.db')
                        
                        command.upgrade(alembic_cfg, 'head')
                        print('列添加成功')
                except Exception as e:
                    print(f'检查表结构时出错: {e}')
    else:
        print('数据库已初始化，跳过初始化步骤')
    
    app.run(debug=True, host='127.0.0.1', port=5000)
